import numpy as np
from typing import Any, Union, Optional, Callable

from tianshou.data import Batch, to_numpy

from agent.policy.dqn import CustomDQNPolicy
from ril.policy.rule_policy_base import RulePolicyBase


class DQNRulePolicy(CustomDQNPolicy, RulePolicyBase):

    def __init__(
        self,
        rule_fn: Callable = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        CustomDQNPolicy.__init__(self, *args, **kwargs)
        RulePolicyBase.__init__(self, rule_fn)

    def forward(
        self,
        batch: Batch,
        state: Optional[Union[dict, Batch, np.ndarray]] = None,
        model: str = "model",
        input: str = "obs",
        **kwargs: Any,
    ) -> Batch:
        model = getattr(self, model)
        obs = batch[input]
        obs_ = obs.obs if hasattr(obs, "obs") else obs
        logits, h = model(obs_, state=state, info=batch.info)
        q = self.compute_q_value(logits)
        if not hasattr(self, "max_action_num"):
            self.max_action_num = q.shape[1]
        act: np.ndarray = to_numpy(q.max(dim=1)[1])

        if hasattr(obs, "mask"):
            # some of actions are masked, they cannot be selected
            q_: np.ndarray = to_numpy(q)
            q_[~obs.mask] = -np.inf
            act = q_.argmax(axis=1)

        # use ril
        if not self.updating:
            act, mask, explorable = self.process_rule(batch, act)
            # 将 mask, explorable 加入 collector 的 data
            # 如果规则未生效，mask = [None, ...]，explorable = [True, ...]
            batch.update(mask=mask, explorable=explorable)

        if hasattr(batch, "mask") and not np.any(batch.mask == None):
            # this mask come from ril
            q_: np.ndarray = to_numpy(q)
            q_[~batch.mask] = -np.inf
            act = q_.argmax(axis=1)

        return Batch(logits=logits, act=act, state=h)

    def exploration_noise(self, act: np.ndarray, batch: Batch) -> np.ndarray:
        if not np.isclose(self.eps, 0.0):
            for i in range(len(act)):
                if np.random.rand() < self.eps and batch.explorable[i]:
                    q_ = np.random.rand(self.max_action_num)
                    if hasattr(batch["obs"], "mask"):
                        q_[~batch["obs"].mask[i]] = -np.inf
                    if hasattr(batch, "mask") and batch.mask[i] is not None:
                        q_[~batch.mask[i]] = -np.inf
                    act[i] = q_.argmax()
        return act
